#!/usr/bin/python
#***************************************************************************
# Copyright notice for this file:
#  Copyright (C) 2012 TheAxeMan
#  _______ _                               _    _
# |__   __| |     ___     /\          ___ | \  / |
#    | |  | |__  / _ \   /  \  __  _ / _ \|  \/  | __ _ _ __
#    | |  | '_ \|  __/  / /\ \ \ \/ |  __/| .  . |/ _` | '_ \
#    | |  | | | | (__/\/ ____ \ )  (| (__/| |\/| | (_| | | | |
#    |_|  |_| |_|\____/_/    \_._/\_.\____|_|  |_|\__,_|_| |_|
#
# Final Fantasy TAS scripts by TheAxeMan
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
#
#***************************************************************************


import sys
import struct
import string

from battlePlannerLib import BattleRandGenerator

#This script needs to get some information from the rom. Set this path
#to point to a Final Fantasy rom.
rom_file_name = "./FinalFantasy-TAS.nes"

#enumeration for battle setup types
class BattleSetups:
  allSmall = 0
  allLarge = 1
  mix = 2

#Hash to store different enemy arrangments. Unless you kill
#some enemies, there is a pretty limited variation, making
#this much easier.
#useful members:
# hashes - list structure with all hashes
# totalEnemies - list indexed by hash to give total enemies
# numHashes - number of hashes (not the max hash as hashes start at 0)
class BattleHash:
  def __init__(self):
    #nine combinations for allSmall setup
    allSmallHashes = [0 for x in range(9)]
    #four combinations for allLarge setup
    allLargeHashes = [0 for x in range(4)]
    #20 combinations for mix setup (3*7-1), set in two dimensions
    #note that there could be no smalls or no larges
    mixHashes = [[0 for x in range(7)] for y in range(3)]

    #also make a list that tracks the number of enemies for each hash
    self.totalEnemies = []

    i = 0
    for x in range(9):
      allSmallHashes[x] = i
      self.totalEnemies.append(x+1)
      i += 1
    for x in range(4):
      allLargeHashes[x] = i
      self.totalEnemies.append(x+1)
      i += 1
    for x in range(3):
      for y in range(7):
        if x == 0 and y == 0: continue
        mixHashes[x][y] = i
        self.totalEnemies.append(x + y)
        i += 1
    self.numHashes = i

    #combine into a list
    self.hashes = [0 for x in range(3)]
    self.hashes[BattleSetups.allSmall] = allSmallHashes
    self.hashes[BattleSetups.allLarge] = allLargeHashes
    self.hashes[BattleSetups.mix] = mixHashes

  #TODO: description is completely wrong....
  #Create a data structure that contains a boolean that is true
  #if that rand allows the character to run. Follows my convention
  #of recording the seed value on the frame when A button is hit.
  #Only considers cases where the player goes first and runs
  #successfully. Of course, in a preemptive battle, this isn't
  #necessary to calculate as the player always goes first and runs
  #successfully.
  def getRunData(self, luck, num_alive=1):
    rg = BattleRandGenerator()
    run_data = []
    for h in range(self.numHashes):
      enemies_present = [False for x in range(9)]
      (setup, num_large, num_small) = self.getInfo(h)

      #first we need to assemble a list that details which
      #enemy slots are occupied
      if setup == BattleSetups.allSmall:
        for x in range(num_small):
          enemies_present[x] = True
      elif setup == BattleSetups.allLarge:
        for x in range(num_large):
          enemies_present[x] = True
      else:   #assuming mix here...
        for x in range(num_large):
          enemies_present[x] = True
        for x in range(num_small):
          enemies_present[x+2] = True

      #now calculate for all seeds
      run_seeds = []
      for i in range(256):
        rg.seed = i
        #find out if player goes first
        order = rg.getBattleOrder()
        party_first = False
        for j in order:
          if j == 80:
            party_first = True
            break
          elif num_alive >= 2 and j == 81:
            party_first = True
            break
          elif j >= 80: continue
          #this has to be last since 80+ will be an invalid index
          if enemies_present[j]: break
        if not party_first: continue
        #now see if the rg allows running
        roll = rg.getRandInRange(0, 16)
        if luck - roll <= 0: continue
        #we ran, so add to list
        run_seeds.append(i)

      run_data.append(run_seeds)

    return run_data

  #Create a list indexed by random seed that pre-calculates the
  #party's surprise roll for each rand. This number gets compared
  #with the enemy group's surprise factor to determine surprise
  #or premptive:
  #party roll - surprise factor <= 10 : surprise!
  #party roll - surprise factor >= 90 : preemptive!
  #note: Indexed by seed used for calculation, not the seed
  #      set before counting enemies.
  def createSurpriseData(self, run_index):
    party_rolls = []
    rg = BattleRandGenerator()
    rg.seed = 0
    for i in range(256):
      roll = run_index + rg.getRandInRange(run_index, 100)
      party_rolls.append(roll)

    return party_rolls

  #Given the setup and number of small and/or large enemies,
  #return the battle hash. The index into the hash table starts at
  #0, so subtract one from enemy total except for mix setup.
  def getHash(self, setup, numLarge, numSmall):
    if setup == BattleSetups.allSmall:
      if numSmall > 9: numSmall = 9
      return self.hashes[setup][numSmall-1]

    elif setup == BattleSetups.allLarge:
      if numLarge > 4: numLarge = 4
      return self.hashes[setup][numLarge-1]

    elif setup == BattleSetups.mix:
      if numLarge > 2: numLarge = 2
      if numSmall > 6: numSmall = 6
      #note that there might be 0 large or small, so don't subtract 1
      return self.hashes[setup][numLarge][numSmall]
    else:
      return 0

  #returns (setup, numLarge, numSmall)
  def getInfo(self, h):
    for x in range(9):
      if h == self.hashes[0][x]: return (BattleSetups.allSmall, 0, x+1)
    for x in range(4):
      if h == self.hashes[1][x]: return (BattleSetups.allLarge, x+1, 0)
    for x in range(3):
      for y in range(7):
        if x == 0 and y == 0: continue
        if h == self.hashes[2][x][y]: return (BattleSetups.mix, x, y)
    #return -1s to show failure...
    return (-1, -1, -1)

  #Compare two hashes. -1 if h1 is strictly worse than h2
  #1 if h2 is strictly worse than h2. 0 if not comparable
  def compare(self, h1, h2):
    return 0


#one battle group
#Members:
# minEnemies - list of minimum number enemies for all four slots
# maxEnemies - list of max number enemies for all four slots
# enemyIsLarge  - list of booleans for all four slots, True means large
# surpriseFactor - surprise factor for the battle
# setup - setup from BattleSetups enum class
class BattleGroup:
  def __init__(self):
    self.minEnemies = [0 for x in range(4)]
    self.maxEnemies = [0 for x in range(4)]
    self.enemyIsLarge = [False for x in range(4)]
    self.surpriseFactor = 4
    self.setup = BattleSetups.allSmall

  #For the group data, calculate the battle hash resulting
  #from each starting seed.
  def calcBattleHashes(self):
    self.battleHashes = []
    rg = BattleRandGenerator()
    bh = BattleHash()
    
    self.seedsUsed = [4 for i in xrange(256)]

    for i in range(256):
      rg.seed = i
      numSmall = 0
      numLarge = 0
      seedsUsed = 0
      gotSixSmall = False
      for j in range(4):
        numEnemies = rg.getRandInRange(self.minEnemies[j], self.maxEnemies[j])
        if self.enemyIsLarge[j]:
          numLarge += numEnemies
          seedsUsed += 1
        else:
          numSmall += numEnemies
          if not gotSixSmall:
            seedsUsed += 1
          if numSmall >= 6:
            gotSixSmall = True
          
      if self.setup == BattleSetups.mix:
        #if seedsUsed < 4:
          #seedsUsed += 1
        self.seedsUsed[i] = seedsUsed
        
      hashResult = bh.getHash(self.setup, numLarge, numSmall)
      self.battleHashes.append(hashResult)
      
  #get the minimum total enemies that can appear in this fight
  def minTotal(self):
    total = 0
    for i in range(4):
      total += self.minEnemies[i]
    return total

  #print out group data
  def printGroup(self):
    for i in range(4):
      if self.maxEnemies[i] == 0: continue
      if self.enemyIsLarge[i]:
        e_size = "large"
      else:
        e_size = "small"
      num_desc = str(self.minEnemies[i]) + "-" + str(self.maxEnemies[i])
      print "Enemy",i,"is",e_size,"and there are",num_desc,"of them"
    if self.setup == BattleSetups.allSmall:
      print "setup is all small"
    elif self.setup == BattleSetups.allLarge:
      print "setup is all large"
    elif self.setup == BattleSetups.mix:
      print "setup is mix"
    else:
      print "setup is unknown"
    print "surprise factor is",self.surpriseFactor



#This class will hold data on battles: enemy count, formation, etc.
class BattleGroups:
  def __init__(self):
    self.loadData()

  def loadData(self):
    self.groups = []
    #This should be a Final Fantasy ROM
    rom_file = open(rom_file_name, "rb")
    read_pos = 0x2C410

    #for group in range(12):
    for group in range(128):
      #each group has 16 bytes of data
      data_list = []
      rom_file.seek(read_pos + group*16)
      for i in range(16):
        data_list.append( ord(rom_file.read(1)) )

      bg = BattleGroup()
      bg.setup = data_list[0] / 16
      bg.surpriseFactor = data_list[12]

      #don't worry about non-standard fights
      if bg.setup > 2:
              #insert a None into the list as a placeholder
        self.groups.append(None)
        continue

      for i in range(4):
        bg.minEnemies[i] = data_list[6+i] / 16
        bg.maxEnemies[i] = data_list[6+i] % 16
        bg.enemyIsLarge[i] = ( (data_list[1] / (2**(2*i))) % 2 == 1 )

      #print "\ngroup number",group
      #bg.printGroup()
      bg.calcBattleHashes()
      self.groups.append(bg)

      #Now repeat for formation 2. All data is in the same 16 bytes.
      #But there are only two enemy types and the numbers are different.

      bg = BattleGroup()
      bg.setup = data_list[0] / 16
      bg.surpriseFactor = data_list[12]

      for i in range(2):
        bg.minEnemies[i] = data_list[14+i] / 16
        bg.maxEnemies[i] = data_list[14+i] % 16
      #The enemyIsLarge data is needed to get seed count
      for i in range(4):
        bg.enemyIsLarge[i] = ( (data_list[1] / (2**(2*i))) % 2 == 1 )

      #print "\ngroup number",group,"formation 2"
      #bg.printGroup()
      bg.calcBattleHashes()
      self.groups.append(bg)

    rom_file.close()

  #Get a battle from the list using name. Format of name is two
  #hex digits and "a" or "b".
  def getBattleByName(self, bname):
    #first chop off the a or b
    if len(bname) < 2: return None
    ab_part = bname[-1:]
    num_part = bname[:-1]

    #now calculate index
    if ab_part == "a":
      index = 2 * int("0x"+num_part, 16)
    elif ab_part == "b":
      index = 2 * int("0x"+num_part, 16) + 1
    else:
      print "invalid battle name"
      sys.exit(0)
    return self.groups[index]

#compares the results from BattleHash::totalEnemies to bh::getInfo
def checkTotals():
  bgs = BattleGroups()
  bh = BattleHash()

  for bg in bgs.groups:
    for x in bg.battleHashes:
      (setup, num_large, num_small) = bh.getInfo(x)
      total = bh.totalEnemies[x]
      if num_large+num_small != total:
        print "error: bad info or total for hash",x
        bg.printGroup()
      #print "for seed",x,"there are",num_small,"enemies"
      #print "for seed",x,"there are",total,"enemies"

def printSurpriseStats():
  bgs = BattleGroups()
  bh = BattleHash()

  s_data = bh.createSurpriseData(2)

  pre_seeds = []
  surprise_seeds = []
  initiative = 90
  for i in range(len(s_data)):
    if s_data[i]-initiative >= 90:
      pre_seeds.append(i)
    elif s_data[i]-initiative <= 10:
      surprise_seeds.append(i)

  print len(pre_seeds),"seeds result in a preemptive battle:",len(pre_seeds)/2.56,"percent"
  print pre_seeds
  print len(surprise_seeds),"seeds result in a surprise battle:",len(surprise_seeds)/2.56,"percent"
  print surprise_seeds
  
  normal_seeds = [i for i in xrange(256) if i not in (pre_seeds+surprise_seeds)]
  print len(normal_seeds),"seeds result in a normal battle",len(normal_seeds)/2.56,"percent"
  print normal_seeds

#for testing, prints out some stats on running
def printRunStats():
  bh = BattleHash()

  report_file = open("bh_report.txt", "w")
  #report_file = sys.stdout

  for luck in range(5,18):
    rd = bh.getRunData(luck)
    print >> report_file, "this data is for",luck,"luck"
    total_run_seeds = 0
    for h in range(len(rd)):
      (setup, num_large, num_small) = bh.getInfo(h)
      num_run_seeds = len(rd[h])
      total_run_seeds += num_run_seeds
      run_percent = "%(#)2.2f" % {"#": num_run_seeds * 100 / 256.0}
      run_percent += "%"
      print >> report_file, "For hash",h,"there are",num_large+num_small,"enemies."
      print >> report_file, "  The chances of running are",num_run_seeds,"in 256 or",run_percent
      print >> report_file

    print >> report_file, "total is",total_run_seeds
    print >> report_file

  if report_file != sys.stdout:
    report_file.close()

def reportRunDelays():
  bg = BattleGroups()
  bh = BattleHash()
  rg = BattleRandGenerator()
  rg.seed = 14

  b = bg.getBattleByName("Ca")


  #generate surprise data
  run_index = 2
  party_rolls = bh.createSurpriseData(run_index)
  #generate run data
  luck = 9
  run_data = bh.getRunData(luck)

  #figure out the corresponding hash
  h = b.battleHashes[rg.seed]
  total_enemies = bh.totalEnemies[h]
  print "met",total_enemies,"enemies"

  #advance to the right rand seed
  rg.useUpRands(4)

  #check surprise status
  surprise_roll = party_rolls[rg.seed]
  if surprise_roll-b.surpriseFactor >= 90: print "preemptive!"
  if surprise_roll-b.surpriseFactor <= 10: print "surprise!"

  rg.useUpRands(3)

  #now go through and report how many frames to wait and run
  run_seeds = run_data[h]
  waits = []
  for rs in run_seeds:
    wait_seeds = rs - rg.seed
    if wait_seeds < 0: wait_seeds += 256
    waits.append( 42240 + 2*wait_seeds )
    #waits.append( 2*wait_seeds )
  waits.sort()
  print waits


#given the run data, a battle hash and the current seed,
#figure out how long a delay is necessary.
def getMinRunDelay(run_data, bh, seed):
  run_seeds = run_data[bh]
  #error if no chance to run
  if len(run_seeds) == 0:
    print "unable to run!!"
    sys.exit(0)
  #figure out the first runnable seed
  running_seed = run_seeds[0]
  for rs in run_seeds:
    if rs >= seed:
      running_seed = rs
      break
  #calculate delay
  delay = running_seed - seed
  if delay < 0: delay += 256
  return delay

def main():
  bg = BattleGroups()
  bh = BattleHash()
  #name_list = ["0a", "7a", "Bb", "Ca", "12a", "1Cb"]
  #this goes from pirates to wizards
  #name_list = "5Cb,12a,7b,7b,Ba,1b,4b,66a,10a".split(",")
  #here is Astos on for a while
  name_list = "Da,Ca,12a,5Ba,0a,5Ba,Ba,5Ba,5Ca,12a,Ba,Da".split(",")

  for i in range(0):
    name_list += name_list[:]
  print len(name_list),"enemies in group"

  #make a list of battle groups
  battles = []
  for name in name_list:
    b = bg.getBattleByName(name)
    battles.append(b)
    #b.printGroup()
    #print

  #now go through and report how to run
  luck = 8
  for luck in range(5,16):
    run_index = 2
    party_rolls = bh.createSurpriseData(run_index)
    rd = bh.getRunData(luck)

    #just get the first runnable seed. Don't worry about surprise, but report it.
    rg = BattleRandGenerator()
    total_wait = 0
    for b in battles:
      #get battle hash
      h = b.battleHashes[rg.seed]
      total_enemies = bh.totalEnemies[h]
      #print "met",total_enemies,"enemies"
      #four seeds for the battle determination
      rg.useUpRands(4)
      #check surprise status
      surprise_roll = party_rolls[rg.seed]
      #if surprise_roll-b.surpriseFactor >= 90: print "preemptive!"
      #if surprise_roll-b.surpriseFactor <= 10: print "surprise!"
      #one seed for the surprise roll
      rg.useUpRands(1)
      #two seeds for the minimum delay
      rg.useUpRands(2)
      #now figure out run delay
      delay = getMinRunDelay(rd, h, rg.seed)
      #print "needed to wait",delay*2,"frames"
      total_wait += delay*2
      #use up the delay seeds
      rg.useUpRands(delay)
      #use up 35 for ordering and run probability
      rg.useUpRands(35)

    print "with",luck,"luck total wait was",total_wait

if __name__ == "__main__":
  #main()
  #reportRunDelays()
  printSurpriseStats()

